package main

import (
	"fmt"
	"image"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"slices"
	"strings"
	"sync"

	// "time"

	"image/jpeg"
	"image/png"

	"github.com/disintegration/imaging"
	"golang.org/x/image/bmp"
	"golang.org/x/image/tiff"
)

var cpu_count = runtime.NumCPU()

const __version__ = "0.1.0"

func main() {
	fmt.Println("Script of automatically converting pictures to JPEG or PNG format.\nVersion ", __version__)
	// Get path of all pictures.
	root_path := "./"
	picture_paths := GetPicturePaths(root_path)
	log.Println("Picture paths:", picture_paths)

	// 多线程转换图片
	picture_path_channel := make(chan string)
	var wg sync.WaitGroup

	// 1. 将图片路径发送到 channel
	wg.Add(1)
	go func() {
		defer wg.Done()
		log.Println("Adding picture paths to channel...")
		for _, picture_path := range picture_paths {
			log.Println("Adding", picture_path)
			picture_path_channel <- picture_path
		}
		close(picture_path_channel)
	}()

	// 启用多线程
	for thread := 0; thread < cpu_count; thread++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for picture_path := range picture_path_channel {
				picture, err := os.Open(picture_path)
				if err != nil {
					log.Printf("Warning: Error Opening %v: %v\n", picture_path, err)
					continue
				}
				defer picture.Close()

				// 2. 从 channel 中取出图片路径，并读取图片
				var img image.Image
				{
					tiff_exts := []string{".tiff", ".tif"}
					jpg_exts := []string{".jpg", ".jpeg"}

					switch {
					case strings.ToLower(filepath.Ext(picture_path)) == ".bmp":
						img, err = bmp.Decode(picture)
					case slices.Contains(tiff_exts, strings.ToLower(filepath.Ext(picture_path))):
						img, err = tiff.Decode(picture)
					case slices.Contains(jpg_exts, strings.ToLower(filepath.Ext(picture_path))):
						img, err = jpeg.Decode(picture)
					case filepath.Ext(picture_path) == ".png":
						img, err = png.Decode(picture)
					default:
						log.Printf("Warning: Unsupported file format for %v\n", picture_path)
						continue
					}
				}

				// 检查是否成功解码
				if img == nil {
					log.Println("Failed to decode", picture_path, "`img` is `nil`!")
					continue
				}
				log.Println("Successfully decoded", picture_path)

				// 设置保存路径
				save_img_name := strings.Replace(picture_path, filepath.Ext(picture_path), "", -1)
				// 设置 PNG 编码器压缩等级
				encoder := png.Encoder{CompressionLevel: png.BestCompression}

				// 3. 若图片长度或宽度大于 500， 则缩放并保存
				newWidth := 500
				newHeight := 500
				originalWidth := img.Bounds().Dx()
				originalHeight := img.Bounds().Dy()
				if float32(originalWidth)/float32(originalHeight) > float32(newWidth)/float32(newHeight) {
					if originalWidth > (newWidth + 10) {
						log.Println("Resampling", picture_path)
						scaleFactor := float64(newWidth) / float64(originalWidth)
						newHeight = int(float64(originalHeight) * scaleFactor)
						resizedImg := imaging.Resize(img, newWidth, newHeight, imaging.Lanczos)
						// 保存调整后尺寸的图像
						resampled_save_img_path := save_img_name + "_resampled.png"
						resampled_save_img, err := os.Create(resampled_save_img_path + "_temp")
						if err != nil {
							log.Printf("WARNING: Error when creating new file `%v`: %v", resampled_save_img_path+"_temp", err)
							continue
						}
						err = encoder.Encode(resampled_save_img, resizedImg)
						if err != nil {
							resampled_save_img.Close()
							log.Printf("Warning: Error when encoding %v as .png: %v\n", resampled_save_img_path, err)
							continue
						}
						err = resampled_save_img.Close()
						if err != nil {
							defer resampled_save_img.Close()
							log.Printf("Error: Error when closing %v: %v\n", resampled_save_img_path, err)
						}
						err = os.Rename(resampled_save_img_path+"_temp", resampled_save_img_path)
						if err != nil {
							log.Panicf("Error: Error when renaming %v: %v\n", resampled_save_img_path, err)
						}
					}
				} else {
					if originalHeight > (newHeight + 10) {
						log.Println("Resampling", picture_path)
						scaleFactor := float64(newHeight) / float64(originalHeight)
						newWidth = int(float64(originalWidth) * scaleFactor)
						resizedImg := imaging.Resize(img, newWidth, newHeight, imaging.Lanczos)
						// 保存调整后尺寸的图像
						resampled_save_img_path := save_img_name + "_resampled.png"
						resampled_save_img, err := os.Create(resampled_save_img_path + "_temp")
						if err != nil {
							log.Printf("WARNING: Error when creating new file `%v`: %v", resampled_save_img_path+"_temp", err)
							continue
						}
						err = encoder.Encode(resampled_save_img, resizedImg)
						if err != nil {
							resampled_save_img.Close()
							log.Printf("Warning: Error when encoding %v as .png: %v\n", resampled_save_img_path, err)
							err = os.Remove(resampled_save_img_path + "_temp")
							if err != nil {
								log.Printf("Error: Error when removing temp file after failing to encode %v: %v\n", resampled_save_img_path+"_temp", err)
								continue
							}
							continue
						}
						resampled_save_img.Close()
						err = os.Rename(resampled_save_img_path+"_temp", resampled_save_img_path)
						if err != nil {
							log.Panicf("Error: Error when renaming %v: %v\n", resampled_save_img_path, err)
						}
					}
				}

				// 4. 如果是TIF，BMP格式图片，将图片转换为 PNG 格式
				if filepath.Ext(picture_path) == ".bmp" || filepath.Ext(picture_path) == ".tiff" || filepath.Ext(picture_path) == ".tif" {
					save_img_path := save_img_name + ".png"
					save_img, err := os.Create(save_img_path + "_temp")
					if err != nil {
						log.Printf("WARNING: Error when creating new file `%v`: %v", save_img_path, err)
						continue
					}
					log.Println("Converting", picture_path)
					err = encoder.Encode(save_img, img)
					if err != nil {
						save_img.Close()
						log.Printf("Error when saving `%v`: %v", save_img_path+"_temp", err)
						continue
					}
					save_img.Close()
					err = os.Rename(save_img_path+"_temp", save_img_path)
					if err != nil {
						err = os.Remove(save_img_path + "_temp")
						if err != nil {
							log.Printf("Error: Error when removing %v: %v\n", save_img_path+"_temp", err)
						}
						log.Printf("Error: Error when renaming %v: %v\n", save_img_path, err)
					}
				}
			}

		}()
	}

	// 等待所有 goroutine 完成
	wg.Wait()
	fmt.Println("Done.")
}

func GetPicturePaths(root string) []string {
	var paths []string
	var filtered_paths []string
	var image_exts []string = []string{
		".bmp", ".tiff", ".tif", ".jpg", ".jpeg", ".png",
	}
	// var processed map[string]bool

	// WalkFunc 用于遍历目录树
	var WalkFunc = func(path string, info os.FileInfo, err error) error {
		if err != nil {
			fmt.Println("Error:", err)
			return err
		}
		if !info.IsDir() {
			ext := strings.ToLower(filepath.Ext(info.Name()))
			if slices.Contains(image_exts, ext) {
				paths = append(paths, path)
			}
		}
		return nil
	}

	// 开始遍历目录
	err := filepath.Walk(root, WalkFunc)
	if err != nil {
		fmt.Println("Failed to walk the directory:", err)
		return nil
	}

	// 如果目录中存在 PNG，不将同名 BMP、TIFF 转换为 PNG，也不会进行图片缩放。
	for _, path := range paths {
		log.Println("Get: ", path)
		if filepath.Ext(path) == ".bmp" {
			name := path[:len(path)-len(filepath.Ext(path))]
			skipFlag := false
			for _, p := range paths {
				if strings.Contains(p, name) && (filepath.Ext(p) == ".png" || filepath.Ext(p) == ".jpg" || filepath.Ext(p) == ".jpeg") {
					log.Printf("Skip %v\n", p)
					skipFlag = true
					break
				}
			}
			if !skipFlag {
				filtered_paths = append(filtered_paths, path)
			}
		} else {
			filtered_paths = append(filtered_paths, path)
		}
	}
	log.Println("Filtered: ", filtered_paths)
	return filtered_paths
}
